import numpy as np

from .wheel_controller import WheelController


class MecanumWheelController(WheelController):
    def __init__(self, r, l, w):
        super().__init__()
        self._r = r
        self._l = l
        self._w = w

        self._control_matrix = np.array([
            [-self._l - self._w, 1, -1],
            [self._l + self._w, 1, 1],
            [self._l + self._w, 1, -1],
            [-self._l - self._w, 1, 1],
        ]) / self._r

    def ctrl(self, vx, vy, wz):
        return self._control_matrix @ np.array([wz, vx, vy])
